/*
 * Copyright 2021 <copyright holder> <email>
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <g2o/core/base_vertex.h>
#include <g2o/core/base_binary_edge.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/solvers/csparse/linear_solver_csparse.h>
#include <g2o/core/robust_kernel_impl.h>
#include <istream>

#include "common.h"
#include "sophus/se3.hpp"

using namespace Sophus;
using namespace Eigen;
using namespace std;

///姿态与内参构造
struct PoseAndIntrinsics{
    PoseAndIntrinsics(){}
    
    ///set from given data address
    explicit PoseAndIntrinsics(double *data_addr){
        rotation=SO3d::exp(Vector3d(data_addr[0],data_addr[1],data_addr[2]));
        translation=Vector3d(data_addr[3],data_addr[4],data_addr[5]);
        focal=data_addr[6];
        k1=data_addr[7];
        k2=data_addr[8];
    }
    
    ///将估计值放入内存
    void set_to(double *data_addr){
        auto r=rotation.log();
        for(int i=0;i<3;i++) data_addr[i]=r[i];
        for(int i=0;i<3;++i) data_addr[i+3]=translation[i];
        data_addr[6]=focal;
        data_addr[7]=k1;
        data_addr[8]=k2;
    }
    
    SO3d rotation;
    Vector3d translation=Vector3d::Zero();
    double focal=0;
    double k1=0,k2=0;
};

/// 位姿加相机内参的顶点，9维，前三维为so3，接下去为t, f, k1, k2
class VertexPoseAndIntrinsics : public g2o::BaseVertex<9, PoseAndIntrinsics> {
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    VertexPoseAndIntrinsics() {}

    virtual void setToOriginImpl() override {
        _estimate = PoseAndIntrinsics();
    }

    virtual void oplusImpl(const double *update) override {
        _estimate.rotation = SO3d::exp(Vector3d(update[0], update[1], update[2])) * _estimate.rotation;
        _estimate.translation += Vector3d(update[3], update[4], update[5]);
        _estimate.focal += update[6];
        _estimate.k1 += update[7];
        _estimate.k2 += update[8];
    }

    /// 根据估计值投影一个点
    Vector2d project(const Vector3d &point) {
        Vector3d pc = _estimate.rotation * point + _estimate.translation;
        pc = -pc / pc[2];
        double r2 = pc.squaredNorm();
        double distortion = 1.0 + r2 * (_estimate.k1 + _estimate.k2 * r2);
        return Vector2d(_estimate.focal * distortion * pc[0],
                        _estimate.focal * distortion * pc[1]);
    }

    virtual bool read(istream &in) {}

    virtual bool write(ostream &out) const {}
};


class VertexPoint:public g2o::BaseVertex<3,Vector3d> {
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
     
    VertexPoint() {}
    
    virtual void setToOriginImpl() override{
        _estimate=Vector3d(0,0,0);
    }
    
    virtual void oplusImpl(const double *update) override{
        _estimate+=Vector3d(update[0],update[1],update[2]);
    }
    
    virtual bool read(istream &in) {}
    virtual bool write(ostream &out) const {}
};

class EdgeProjection:public g2o::BaseBinaryEdge<2,Vector2d,VertexPoseAndIntrinsics,VertexPoint>{
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
     
    virtual void computeError() override{
        auto v0=(VertexPoseAndIntrinsics *)_vertices[0];
        auto v1=(VertexPoint *) _vertices[1];
        auto proj=v0->project(v1->estimate());
        _error=proj-_measurement;
    }
    
    //use numeric dercatives
    virtual bool read(istream &in) {}
    virtual bool write(ostream &out) const{}
};

void SolveBA(BALProblem &bal_problem);

int main(int argc, char** argv){
    if(argc!=2){
        cout<<"usage: bundle_adjustment_g2o bal_data.txt"<<endl;
        return 1;
    }
    
    BALProblem bal_problem(argv[1]);
    bal_problem.Normalize();
    bal_problem.Perturb(0.1,0.5,0.5);
    bal_problem.WriteToPLYFile("initial.ply");
    SolveBA(bal_problem);
    bal_problem.WriteToPLYFile("final.ply");
    
    return 0;
}

 void SolveBA(BALProblem& bal_problem)
{
    const int point_block_size=bal_problem.point_block_size();
    const int camera_block_size=bal_problem.camera_block_size();
    double *points=bal_problem.mutable_points();
    double *camera=bal_problem.mutable_cameras();
    
    //pose dimension 9, landmark is 3
    typedef g2o::BlockSolve<g2o::BlockSolverTraits<9,3>> BlockSolverType;
    typedef g2o::LinearSolverCSparse<BlockSolverType::PoseMatrixType> LinearSolverType;
    //use LM
    auto solver=new g2o::OptimizationAlgorithmLevenberg(
        g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>()));
    g2o::SpareOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose((true);
    
    ///build g2o bal_problem
    const double *observations=bal_problem.observations();
    //Vertex
    vector<VertexPoseAndIntrinsics *> vertex_pose_intrinsics;
    vector<VertexPoint *> Vertex_points;
    for(int i=0;i<bal_problem.num_cameras();++i){
        double *camera=camera+camera_block_size*i;
        v->setId(i);
        v->setEstimate(PoseAndIntrinsics(camera));
        optimizer.addVertex(v);
        vertex_pose_intrinsics.push_back(v);
    }
    for(int i=0;i<bal_problem.num_points();++i){
        VertexPoint *v=new VertexPoint();
        double *point=point+point_block_size*i;
        v->setId(i+bal_problem.num_cameras());
        v->setEstimate(Vector3d(point[0],point[1],point[2]));
        
        //g2o在BA中手动设置marg的顶点
        v->setMarginalized(true);
        optimizer.addVertex(v);
        Vertex_points.push_back(v);
    }
    
    //Edge
    for(int i=0;i<bal_problem.num_observations();++i)
    {
        EdgeProjection *edge=new EdgeProjection;
        edge->setVertex(0, vertex_pose_intrinsics[bal_problem.camera_index()[i]]);
        edge->setVertex(1,vertex_pose_intrinsics[bal_problem.point_index()[i]]);
        edge->setMeasurement(Vector2d(observations[2*i+0],observations[2*i+1]));
        edge->setInformation(Matrix2d::Identity());
        edge->setRobustKernel(new g2o::RobustKernelHuber());
        optimizer.addEdge(edge);
    }
    
    optimizer.initializeOptimization();
    optimizer.optimize(40);
    
    //set_to bal bal_problem
    for(int i=0;i<bal_problem.num_cameras();++i)
    {
        double *camera=cameras+camera_block_size*i;
        auto vertex=vertex_pose_intrinsics[i];
        auto estimate=vertex->estimate();
        estimate.set_to(camera);
    }
    for(int i=0;i<bal_problem.num_points();++i){
        double *point=points+point_block_size*i;
        auto vertex=Vertex_points[i];
        for(int k=0;k<3;++k) point[k]=vertex->estimate()[K];
    }
}





























    
